import sqlite3
import pandas as pd
from datetime import datetime
import os

class DataManager:
    def __init__(self, db_path="data/culture_wave.db"):
        # 确保数据目录存在
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self.db_path = db_path
        self._init_database()
    
    def _init_database(self):
        """初始化数据库表结构"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建评论数据表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS comments (
                id INTEGER PRIMARY KEY AUTOINCREMENT, -- 改为AUTOINCREMENT
                video_id TEXT,
                text TEXT,
                sentiment REAL,
                likes INTEGER,
                timestamp DATETIME,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP)
            
        ''')
        
        # 创建分析结果表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS analysis_results (
                id INTEGER PRIMARY KEY AUTOINCREMENT,  -- 改为AUTOINCREMENT
                video_id TEXT,
                entropy_value REAL,
                cultural_phase TEXT,
                analysis_time DATETIME,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP)
            
        ''')
        
        conn.commit()
        conn.close()
    
    def save_comments(self, comments_df, video_id):
        """保存评论数据到数据库（增强错误处理）"""
        if comments_df is None or comments_df.empty:
            print("没有数据可保存")
            return False
            
        conn = sqlite3.connect(self.db_path)
        try:
            # 准备数据
            df_to_save = comments_df.copy()
            df_to_save['video_id'] = video_id
            
            # 确保时间格式正确
            if 'time' in df_to_save.columns:
                df_to_save['time'] = pd.to_datetime(df_to_save['time'])
            
            # 重命名列以匹配数据库
            column_mapping = {
                'text': 'text',
                'sentiment': 'sentiment', 
                'like': 'likes',
                'time': 'timestamp'
            }
            
            # 只重命名存在的列
            for src_col, dest_col in column_mapping.items():
                if src_col in df_to_save.columns:
                    df_to_save[dest_col] = df_to_save[src_col]
            
            # 选择需要的列
            db_columns = ['video_id', 'text', 'sentiment', 'likes', 'timestamp']
            available_columns = [col for col in db_columns if col in df_to_save.columns]
            
            df_to_save = df_to_save[available_columns]
            
            # 保存到数据库
            df_to_save.to_sql('comments', conn, if_exists='append', index=False)
            print(f"成功保存 {len(df_to_save)} 条评论数据")
            return True
            
        except Exception as e:
            print(f"保存数据失败: {e}")
            return False
        finally:
            conn.close()
    
    def save_analysis_result(self, video_id, entropy_value, cultural_phase):
        """保存分析结果（新增）"""
        conn = sqlite3.connect(self.db_path)
        try:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO analysis_results (video_id, entropy_value, cultural_phase, analysis_time)
                VALUES (?, ?, ?, ?)
            ''', (video_id, entropy_value, cultural_phase, datetime.now()))
            conn.commit()
            print(f"分析结果已保存: {cultural_phase} (熵值: {entropy_value:.4f})")
            return True
        except Exception as e:
            print(f"保存分析结果失败: {e}")
            return False
        finally:
            conn.close()
    
    def get_comments(self, video_id=None):
        """获取评论数据（新增）"""
        conn = sqlite3.connect(self.db_path)
        try:
            if video_id:
                query = "SELECT * FROM comments WHERE video_id = ? ORDER BY timestamp DESC"
                df = pd.read_sql_query(query, conn, params=(video_id,))
            else:
                df = pd.read_sql_query("SELECT * FROM comments ORDER BY timestamp DESC", conn)
            return df
        except Exception as e:
            print(f"获取评论数据失败: {e}")
            return pd.DataFrame()
        finally:
            conn.close()
    
    def get_analysis_history(self, video_id=None):
        """获取分析历史（新增）"""
        conn = sqlite3.connect(self.db_path)
        try:
            if video_id:
                query = "SELECT * FROM analysis_results WHERE video_id = ? ORDER BY analysis_time DESC"
                df = pd.read_sql_query(query, conn, params=(video_id,))
            else:
                df = pd.read_sql_query("SELECT * FROM analysis_results ORDER BY analysis_time DESC", conn)
            return df
        except Exception as e:
            print(f"获取分析历史失败: {e}")
            return pd.DataFrame()
        finally:
            conn.close()